package com.fowo.api.sys.component;

import com.fowo.api.sys.annotation.DynamicMultiDataSourceTransactional;
import org.aspectj.lang.annotation.*;
import org.springframework.context.ApplicationContext;
import org.springframework.jdbc.datasource.DataSourceTransactionManager;
import org.springframework.stereotype.Component;
import org.springframework.transaction.TransactionDefinition;
import org.springframework.transaction.TransactionStatus;
import org.springframework.transaction.support.AbstractPlatformTransactionManager;
import org.springframework.transaction.support.DefaultTransactionDefinition;

import javax.annotation.Resource;
import javax.sql.DataSource;
import java.util.HashMap;
import java.util.Map;
import java.util.Stack;

/**
 * 动态多源事务切面
 * @author yl_ls
 */
@Component
@Aspect
public class DynamicMultiDataSourceTransactionalAop {
    private static final ThreadLocal<Stack<Map<DataSourceTransactionManager, TransactionStatus>>> THREAD_LOCAL = new ThreadLocal<>();

    @Resource
    private DataSourceTransactionManager dataSourceTransactionManager;

    private static final DefaultTransactionDefinition def = new DefaultTransactionDefinition();
    static {
        def.setReadOnly(false);
        def.setIsolationLevel(TransactionDefinition.ISOLATION_DEFAULT);
        def.setPropagationBehavior(TransactionDefinition.PROPAGATION_REQUIRED);
    }

    @Pointcut("@annotation(com.fowo.api.sys.annotation.DynamicMultiDataSourceTransactional)")
    public void pointcut() {
    }

    @Before("pointcut() && @annotation(transactional)")
    public void before(DynamicMultiDataSourceTransactional transactional) {
        Stack<Map<DataSourceTransactionManager, TransactionStatus>> pairStack = new Stack<>();
        // 先放默认主数据源事务
        TransactionStatus transaction = dataSourceTransactionManager.getTransaction(def);
        Map<DataSourceTransactionManager, TransactionStatus> transactionMap = new HashMap<>();
        transactionMap.put(dataSourceTransactionManager, transaction);
        pairStack.push(transactionMap);
        THREAD_LOCAL.set(pairStack);
    }

    @AfterReturning("pointcut()")
    public void afterReturning() {
        Stack<Map<DataSourceTransactionManager, TransactionStatus>> pairStack = THREAD_LOCAL.get();
        while (pairStack != null && !pairStack.empty()) {
            Map<DataSourceTransactionManager, TransactionStatus> pair = pairStack.pop();
            pair.forEach((k,v) -> {
                if (!v.isCompleted()) {
                    k.commit(v);
                }
            });
        }
        THREAD_LOCAL.remove();
    }

    @AfterThrowing("pointcut()")
    public void afterThrowing() {
        Stack<Map<DataSourceTransactionManager, TransactionStatus>> pairStack = THREAD_LOCAL.get();
        while (pairStack != null && !pairStack.empty()) {
            Map<DataSourceTransactionManager, TransactionStatus> pair = pairStack.pop();
            pair.forEach((k,v) -> {
                if (!v.isCompleted()) {
                    k.rollback(v);
                }
            });
        }
        THREAD_LOCAL.remove();
    }

    public static void addDataSource(DataSource dataSource, boolean focusCreateTrans) {
        Stack<Map<DataSourceTransactionManager, TransactionStatus>> pairStack = THREAD_LOCAL.get();
        if (pairStack == null) {
            throw new RuntimeException("addDataSource invoke with not DynamicMultiDataSourceTransaction Scope.");
        }
        if (!focusCreateTrans && pairStack.stream().anyMatch(m -> m.keySet().stream().anyMatch(kv -> kv.getDataSource() == dataSource))) {
            return;
        }
        DataSourceTransactionManager transactionManager = new DataSourceTransactionManager(dataSource);
        TransactionStatus transactionStatus = transactionManager.getTransaction(def);
        add(transactionManager, transactionStatus);
    }

    public static void add(DataSourceTransactionManager transactionManager, TransactionStatus transactionStatus) {
        THREAD_LOCAL.get().push(Map.of(transactionManager, transactionStatus));
    }
}
